#include <iostream>
#include <memory>
#include <vector>
#include <string>
#include "yolov3_func.hpp"
#include <inference_engine.hpp>
#include "opencv2/opencv.hpp"

using namespace InferenceEngine;

int main()
{        
    cv::VideoCapture capture;
    cv::Mat frame;
    frame = capture.open("/home/yuanpei-chen/桌面/worker-zone-detection.mp4");
    const size_t width  = 1920;
    const size_t height = 1080;

    //------------------------1--------------------------
    Core ie;
    //-------------------------2-------------------------
    CNNNetwork network = ie.ReadNetwork("/opt/intel/openvino/deployment_tools/model_optimizer/frozen_darknet_yolov3_model.xml","/opt/intel/openvino/deployment_tools/model_optimizer/frozen_darknet_yolov3_model.bin");
    //--------------------------3-----------------------
    InputsDataMap inputInfo(network.getInputsInfo());
    InputInfo::Ptr& input = inputInfo.begin()->second;
    auto inputName = inputInfo.begin()->first;

    input->setPrecision(Precision::U8);
    input->getInputData()->setLayout(Layout::NCHW);
    ICNNNetwork::InputShapes inputShapes = network.getInputShapes();
    SizeVector& inSizeVector = inputShapes.begin()->second;
    inSizeVector[0] = 1;  // set batch to 1
    network.reshape(inputShapes);

    OutputsDataMap outputInfo(network.getOutputsInfo());
    for (auto &output : outputInfo) 
    {
            output.second->setPrecision(Precision::FP32);
            output.second->setLayout(Layout::NCHW);
    }
    std::map<std::string, YoloParams> yoloParams;
        if (auto ngraphFunction = network.getFunction()) 
        {
            for (const auto op : ngraphFunction->get_ops()) 
            {
                auto outputLayer = outputInfo.find(op->get_friendly_name());
                if (outputLayer != outputInfo.end()) 
                {
                    auto regionYolo = std::dynamic_pointer_cast<ngraph::op::RegionYolo>(op);
                    if (!regionYolo) 
                    {
                        throw std::runtime_error("Invalid output type: " +
                            std::string(regionYolo->get_type_info().name) + ". RegionYolo expected");
                    }
                    yoloParams[outputLayer->first] = YoloParams(regionYolo);
                }
            }
        }
    
    //---------------------------4--------------------------
    ExecutableNetwork executable_network = ie.LoadNetwork(network, "CPU");
    //---------------------------5---------------------------
    InferRequest::Ptr infer_request = executable_network.CreateInferRequestPtr();
    //---------------------------6---------------------------
    while(capture.read(frame)) 
    {
        Blob::Ptr frameBlob = infer_request->GetBlob(inputName);
        matU8ToBlob<uint8_t>(frame, frameBlob);

        infer_request->Infer();
        const TensorDesc& inputDesc = inputInfo.begin()->second.get()->getTensorDesc();
        std::vector<std::vector<float>> objects;

        // Parsing outputs
        for (auto &output : outputInfo)
        {
            auto output_name = output.first;
            Blob::Ptr blob = infer_request->GetBlob(output_name);
            std::vector<float> anchors = {10.0, 27.0, 23.0, 58.0, 37.0, 13.0,
                                          82.0, 81.0, 169.0, 135.0, 319.0, 344.0};
            unsigned long resized_im_w = getTensorWidth(inputDesc);
            auto side = resized_im_w;
            auto side_square = side * side;
            LockedMemory<const void> blobMapped = as<MemoryBlob>(blob)->rmap();
            const float *output_blob = blobMapped.as<float *>();
            
            for (int j = 0; j < side * side; j++)
            {
                int row = j / side;
                int col = j % side;
                for (int k = 0; k < 3; k++)
                {
                    float scale = output_blob[side_square * k * 85 + side_square * 4 + j];
                    if (scale < 0.5)
                        continue;
                    double x = (col + output_blob[side_square * k * 85 + j + 0 * side_square]) / side * 416;
                    double y = (row + output_blob[side_square * k * 85 + 1 * side_square]) / side * 416;
                    if (side < 20)
                    {    
                        double height = std::exp(output_blob[side_square * k * 85 + 3 * side_square]) * anchors[2 * k + 1];
                        double width = std::exp(output_blob[side_square * k * 85 + 2 * side_square]) * anchors[2 * k];
                    }
                    else
                    {
                        double height = std::exp(output_blob[side_square * k * 85 + 3 * side_square]) * anchors[2 * k + 7];
                        double width = std::exp(output_blob[side_square * k * 85 + 2 * side_square]) * anchors[2 * k + 6];
                    }
                    for (int l = 0; l < 80; l++) 
                    {
                        float prob = scale * output_blob[side_square * k * 85 + 5 * side_square + l];
                        if (prob < 0.5)
                            continue;
                        float xmin = ((x - width / 2) * static_cast<float>(1920) / static_cast<float>(416));
                        float ymin = ((y - height / 2) * static_cast<float>(1080) / static_cast<float>(416));
                        float xmax = (xmin + width * static_cast<float>(1920) / static_cast<float>(416));
                        float ymax = (ymin + height * static_cast<float>(1080) / static_cast<float>(416));
                        float class_id = j;
                        float confidence = prob;
                        std::vector<float> obj = {xmin, ymin, xmax, ymax, class_id, confidence};
                        objects.push_back(obj);
                    }
                }
            }
        }    
//......................................................
        InverseBubbleSort(objects, objects.size());
        //NMS
        for (size_t i = 0; i < objects.size(); ++i) 
            {
            if (objects[i][5] == 0)
                continue;
            for (size_t j = i + 1; j < objects.size(); ++j)
                if (IntersectionOverUnion(objects[i], objects[j]) >= 0.4)
                    objects[j][5] = 0;
            }

        for (size_t i = 0; i < objects.size(); i++) 
        {
            if (objects[i][5] < 0.5)
                continue;
            if (objects[i][5] > 0.5) 
            {
                cv::rectangle(frame, cv::Point2f(static_cast<float>(objects[i][0]), static_cast<float>(objects[i][1])),
                                  cv::Point2f(static_cast<float>(objects[i][2]), static_cast<float>(objects[i][3])), cv::Scalar(0, 255, 0), 2, 8);
            }
        }
        cv::imshow("Detection results", frame);
        cv::waitKey(10);
    }
    capture.release();
    return 0;
}


